"""
图像识别工具类
"""

import time
import pyautogui
from typing import Optional, Tuple, List, Dict
from PIL import Image
import os


class ImageUtils:
    """图像识别工具类"""
    
    def __init__(self, confidence_levels: List[float] = None, max_retries: int = 3):
        """
        初始化图像工具
        
        Args:
            confidence_levels: 置信度级别列表
            max_retries: 最大重试次数
        """
        self.confidence_levels = confidence_levels or [0.8, 0.7, 0.6, 0.5]
        self.max_retries = max_retries
    
    def locate_image_with_retry(self, image_path: str, max_retries: int = None, 
                              confidence_levels: List[float] = None, 
                              region: Optional[Tuple[int, int, int, int]] = None) -> Optional[Tuple[int, int, int, int]]:
        """
        带重试机制的图像识别
        
        Args:
            image_path: 图像文件路径
            max_retries: 最大重试次数
            confidence_levels: 置信度级别列表
            region: 搜索区域 (x, y, width, height)
            
        Returns:
            图像位置或None
        """
        max_retries = max_retries or self.max_retries
        confidence_levels = confidence_levels or self.confidence_levels
        
        # 检查图像文件是否存在
        if not self.check_image_exists(image_path):
            print(f"图像文件不存在: {image_path}")
            return None
        
        for attempt in range(max_retries):
            print(f"尝试识别图像: {image_path} (第{attempt + 1}次)")
            
            for confidence in confidence_levels:
                try:
                    print(f"  置信度: {confidence}")
                    location = pyautogui.locateOnScreen(
                        image_path, 
                        confidence=confidence,
                        region=region
                    )
                    if location:
                        print(f"  ✅ 找到图像! 位置: {location}")
                        return location
                    else:
                        print(f"  ❌ 未找到图像")
                except pyautogui.ImageNotFoundException:
                    print(f"  ❌ 未找到图像")
                except Exception as e:
                    print(f"  ❌ 发生错误: {e}")
            
            if attempt < max_retries - 1:
                print(f"等待1秒后重试...")
                time.sleep(1)
        
        print(f"❌ 图像识别失败: {image_path}")
        return None
    
    def check_image_exists(self, image_path: str) -> bool:
        """
        检查图像文件是否存在
        
        Args:
            image_path: 图像文件路径
            
        Returns:
            是否存在
        """
        return os.path.exists(image_path)
    
    def get_image_info(self, image_path: str) -> Optional[Dict]:
        """
        获取图像信息
        
        Args:
            image_path: 图像文件路径
            
        Returns:
            图像信息字典
        """
        if not self.check_image_exists(image_path):
            return None
        
        try:
            img = Image.open(image_path)
            return {
                'size': img.size,
                'mode': img.mode,
                'format': img.format
            }
        except Exception as e:
            print(f"获取图像信息失败: {e}")
            return None
    
    def take_screenshot(self, filename: str = None) -> str:
        """
        拍摄截图
        
        Args:
            filename: 文件名
            
        Returns:
            截图文件路径
        """
        if filename is None:
            filename = f'screenshot_{int(time.time())}.png'
        
        screenshot = pyautogui.screenshot()
        screenshot.save(filename)
        print(f"📸 截图已保存: {filename}")
        return filename
    
    def test_image_recognition(self, image_path: str, 
                             confidence_levels: List[float] = None) -> bool:
        """
        测试图像识别效果
        
        Args:
            image_path: 图像文件路径
            confidence_levels: 置信度级别列表
            
        Returns:
            是否识别成功
        """
        confidence_levels = confidence_levels or self.confidence_levels
        
        print(f"\n=== 测试图像识别: {image_path} ===")
        
        # 获取图像信息
        info = self.get_image_info(image_path)
        if info:
            print(f"📏 图像尺寸: {info['size']}")
            print(f"🎨 图像模式: {info['mode']}")
        
        # 测试识别
        location = self.locate_image_with_retry(image_path, confidence_levels=confidence_levels)
        return location is not None
    
    def batch_test_images(self, image_files: List[str]) -> Dict[str, bool]:
        """
        批量测试图像识别
        
        Args:
            image_files: 图像文件列表
            
        Returns:
            测试结果字典
        """
        print("🚀 开始批量测试图像识别...")
        print("请确保目标界面已打开并且相关元素可见")
        
        results = {}
        for image_file in image_files:
            results[image_file] = self.test_image_recognition(image_file)
        
        # 输出测试结果
        print("\n" + "="*50)
        print("📊 测试结果汇总:")
        print("="*50)
        
        success_count = 0
        for image_file, success in results.items():
            status = "✅ 成功" if success else "❌ 失败"
            print(f"{image_file:<15} {status}")
            if success:
                success_count += 1
        
        success_rate = success_count / len(image_files) * 100 if image_files else 0
        print(f"\n成功率: {success_count}/{len(image_files)} ({success_rate:.1f}%)")
        
        return results
